library(tidyverse)
library(ggpubr)

# read gg data --------
vrc01_gg <- read_delim('~/snap-align/gg.corrected.txt')
vrc01_rr <- read_delim('~/snap-align/rr.corrected.txt')

vcr01_seq <- 'CAGGTGCAGCTGGTGCAGTCTGGGGCTGAGGTGAAGAAGCCTGGGGCCTCAGTGAAGGTCTCCTGCAAGGCTTCTGGATACACCTTCACCGGCTACTATATGCACTGGGTGCGACAGGCCCCTGGACAAGGGCTTGAGTGGATGGGATGGATCAACCCTAACAGTGGTGGCACAAACTATGCACAGAAGTTTCAGGGCAGGGTCACCATGACCAGGGACACGTCCATCAGCACAGCCTACATGGAGCTGAGCAGGCTGAGATCTGACGACACGGCCGTGTATTACTGTGCGAGAGGAAAAAACAGCGATTACAATTGGGACTTCCAGCACTGGGGCCAGGGCACCCTGGTCACCGTCTCCTCA'

vcr01_seqlist <- vcr01_seq |> str_extract_all('.') |> pluck(1)

vrc01_gg$count |> sum()
vrc01_rr$count |> sum()

vrc01_gg |>
  mutate(offset = v.del + v.end.in.cdr3 - 9) |>
  select(count, contains('mutations.nt'), offset) |>
  unite(all_mut, contains('mutations.nt'),sep = ',',na.rm = TRUE) |>
  filter(all_mut == '') |>
  summarise(sum(count))

vrc01_rr |>
  mutate(offset = v.del + v.end.in.cdr3 - 9) |>
  select(count, contains('mutations.nt'), offset) |>
  unite(all_mut, contains('mutations.nt'),sep = ',',na.rm = TRUE) |>
  filter(all_mut == '') |>
  summarise(sum(count))

vrc01_gg |> count(v,d,j, sort = TRUE)

vrc01_gg |>
  count(has.cdr3, in.frame, no.stop, complete, canonical)

vrc01_gg |>
  mutate(offset = v.del + v.end.in.cdr3 - 9) |>
  filter(str_detect(mutations.nt.CDR3, 'S286:G>T|S297:A>G')) |>
  head() |>
  DT::datatable()

cdr3 <- vrc01_gg |>
  head(2)

vrc01_gg |>
  filter( !is.na(mutations.nt.CDR2)) |>
  head(2) |>
  bind_rows(cdr3) |>
  DT::datatable()

# only cdr3 sum -------
cdr3_gg <- vrc01_gg |>
  mutate(offset = v.del + v.end.in.cdr3 - 9) |>
  select(freq, mutations.nt.CDR3, offset) |>
  separate_longer_delim(mutations.nt.CDR3, delim = ',') |>
  filter(!is.na(mutations.nt.CDR3)) |>
  mutate(type = str_extract(mutations.nt.CDR3, 'S|D|I'),
         refalt = str_extract(mutations.nt.CDR3, '(?<=:).+'),
         position = str_extract(mutations.nt.CDR3, '\\d+') |> as.numeric(),
         position = position + offset) |>
  count(type,refalt,position, wt = freq) |>
  mutate(ref = str_extract(refalt,'.')) |>
  filter(ref == vcr01_seqlist[position])

cdr3_gg |>
  filter(type == 'S') |>
  ggplot(aes(position, n, fill = refalt)) +
  geom_col() +
  theme_pubr() +
  labs_pubr() +
  geom_vline(xintercept = c(289,330), linetype = 'dashed', color = 'red') +
  expand_limits(y = .055) +
  labs(title = 'GG genotype VRC01 point mutation',
       y = 'Fraction in all sequences')

last_plot() -> g1

cdr3_rr <- vrc01_rr |>
  mutate(offset = v.del + v.end.in.cdr3 - 9) |>
  select(freq, mutations.nt.CDR3, offset) |>
  separate_longer_delim(mutations.nt.CDR3, delim = ',') |>
  filter(!is.na(mutations.nt.CDR3)) |>
  mutate(type = str_extract(mutations.nt.CDR3, 'S|D|I'),
         refalt = str_extract(mutations.nt.CDR3, '(?<=:).+'),
         position = str_extract(mutations.nt.CDR3, '\\d+') |> as.numeric(),
         position = position + offset) |>
  count(type,refalt,position, wt = freq) |>
  mutate(ref = str_extract(refalt,'.')) |>
  filter(ref == vcr01_seqlist[position])

cdr3_rr |>
  filter(type == 'S') |>
  ggplot(aes(position, n, fill = refalt)) +
  geom_col() +
  theme_pubr() +
  labs_pubr() +
  geom_vline(xintercept = c(289,330), linetype = 'dashed', color = 'red') +
  expand_limits(y = .055) +
  labs(title = 'RR genotype VRC01 point mutation',
       y = 'Fraction in all sequences')

g2 <- last_plot()

g1 + g2 +
  patchwork::plot_layout(guides = 'collect') +
  patchwork::plot_annotation(theme = theme(legend.position = 'top'),
                             title = 'CDR3 region is highlighted in red')

cdr3_gg |>
  filter(type == 'I') |>
  summarise(sum(n))

cdr3_rr |>
  filter(type == 'I') |>
  summarise(sum(n))

# full length -------
full_gg <- vrc01_gg |>
  mutate(offset = v.del + v.end.in.cdr3 - 9) |>
  select(freq, contains('mutations.nt'), offset) |>
  unite(all_mut, contains('mutations.nt'),sep = ',',na.rm = TRUE) |>
  separate_longer_delim(all_mut, delim = ',') |>
  filter(!is.na(all_mut)) |>
  mutate(type = str_extract(all_mut, 'S|D|I'),
         refalt = str_extract(all_mut, '(?<=:).+'),
         position = str_extract(all_mut, '\\d+') |> as.numeric(),
         position = position + offset) |>
  count(type,refalt,position, wt = freq) |>
  mutate(ref = str_extract(refalt,'.')) |>
  filter(ref == vcr01_seqlist[position])

full_gg |>
  filter(type == 'S') |>
  ggplot(aes(position, n, fill = refalt)) +
  geom_col() +
  theme_pubr() +
  labs_pubr() +
  geom_vline(xintercept = c(289,330), linetype = 'dashed') +
  geom_label(label = 'CDR3', x = 297, y = 0.032, size = 6, inherit.aes = FALSE)

full_rr <- vrc01_rr |>
  mutate(offset = v.del + v.end.in.cdr3 - 9) |>
  select(freq, contains('mutations.nt'), offset) |>
  unite(all_mut, contains('mutations.nt'),sep = ',',na.rm = TRUE) |>
  separate_longer_delim(all_mut, delim = ',') |>
  filter(!is.na(all_mut)) |>
  mutate(type = str_extract(all_mut, 'S|D|I'),
         refalt = str_extract(all_mut, '(?<=:).+'),
         position = str_extract(all_mut, '\\d+') |> as.numeric(),
         position = position + offset) |>
  count(type,refalt,position, wt = freq) |>
  mutate(ref = str_extract(refalt,'.')) |>
  filter(ref == vcr01_seqlist[position])

full_rr |>
  filter(type == 'S') |>
  ggplot(aes(position, n, fill = refalt)) +
  geom_col() +
  theme_pubr() +
  labs_pubr() +
  geom_vline(xintercept = c(289,330), linetype = 'dashed') +
  geom_label(label = 'CDR3', x = 297, y = 0.032, size = 6, inherit.aes = FALSE) +
  scale_fill_viridis_d(option = 'turbo', begin = .1)

full_rr |>
  filter(type == 'D') |>
  ggplot(aes(position, n, fill = refalt)) +
  geom_col() +
  theme_pubr() +
  labs_pubr() +
  geom_vline(xintercept = c(289,330), linetype = 'dashed') +
  geom_label(label = 'CDR3', x = 297, y = 0.032, size = 6, inherit.aes = FALSE) +
  scale_fill_viridis_d(option = 'turbo', begin = .1)

ss |>
  separate_longer_delim(mutations.nt.CDR3, delim = ',') |>
  filter(!is.na(mutations.nt.CDR3)) |>
  mutate(type = str_extract(mutations.nt.CDR3, 'S|D|I'),
         refalt = str_extract(mutations.nt.CDR3, '(?<=:).+'),
         position = str_extract(mutations.nt.CDR3, '\\d+') |> as.numeric(),
         position = position + offset) |>
  count(type,refalt,position, wt = freq)


sum_gg <- vrc01_gg |>
  select(count, contains('mutations.nt')) |>
  unite(all_mut, contains('mutations.nt'),sep = ',',na.rm = TRUE) |>
  count(all_mut, wt = count) |>
  separate_longer_delim(all_mut, delim = ',') |>
  count(all_mut, wt = n, sort = TRUE) |>
  separate_wider_delim(all_mut, delim = ':', names = c('position','refalt'), too_few = 'align_start') |>
  mutate(type = str_extract(position, 'S|D|I'),
         position = str_remove(position, 'S|D|I') |> as.numeric())

sum_gg |>
  filter(type == 'S') |>
  mutate(n = n/all_read_count) |>
  ggplot(aes(position, n, fill = refalt)) +
  geom_col()

sum_gg |> slice_max(position)

sum_gg <- vrc01_gg |>
  select(count, mutations.nt.CDR3, cdr3nt, contignt) |>
  count(all_mut, wt = count) |>
  separate_longer_delim(all_mut, delim = ',') |>
  count(all_mut, wt = n, sort = TRUE) |>
  separate_wider_delim(all_mut, delim = ':', names = c('position','refalt'), too_few = 'align_start') |>
  mutate(type = str_extract(position, 'S|D|I'),
         position = str_remove(position, 'S|D|I') |> as.numeric())

vrc01_gg |>
  select(count, mutations.nt.CDR3, contignt) |>
  pull(contignt) |>
  pluck(1)

vrc01_rr |>
  select(count, mutations.nt.CDR3, contignt)

